package git.soulbgm;

import git.soulbgm.config.ProxyConfig;
import git.soulbgm.health.BackendHealthChecker;
import git.soulbgm.loadbalancer.RoundRobinLoadBalancer;
import git.soulbgm.stats.TrafficStats;
import git.soulbgm.thread.ThreadPoolManager;
import io.netty.bootstrap.Bootstrap;
import io.netty.buffer.ByteBuf;
import io.netty.buffer.ByteBufUtil;
import io.netty.buffer.Unpooled;
import io.netty.channel.*;
import io.netty.channel.nio.NioEventLoopGroup;
import io.netty.channel.socket.DatagramChannel;
import io.netty.channel.socket.DatagramPacket;
import io.netty.channel.socket.InternetProtocolFamily;
import io.netty.channel.socket.nio.NioDatagramChannel;
import io.netty.util.NetUtil;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import git.soulbgm.pojo.Backend;

import java.net.*;
import java.nio.file.*;
import java.util.Enumeration;
import java.util.concurrent.Future;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;

/**
 * UDP代理服务器的主类
 *
 * <p>该类负责处理UDP数据包的接收和转发功能。它作为客户端和后端服务器之间的中间层，
 * 接收来自客户端的UDP数据包，并通过负载均衡器将其转发到适当的后端服务器。</p>
 *
 * <p>主要功能：</p>
 * <ul>
 *   <li>监听指定端口接收UDP数据包</li>
 *   <li>维护客户端会话信息</li>
 *   <li>通过负载均衡器转发数据包</li>
 *   <li>处理后端服务器响应</li>
 * </ul>
 *
 * @author soulbgm
 * @version 1.0
 * @since 1.0
 */
public class UdpProxyServer {
    /**
     * 日志记录器
     */
    private static final Logger logger = LoggerFactory.getLogger(UdpProxyServer.class);
    /**
     * 服务器运行状态标志
     */
    private static final AtomicBoolean RUNNING = new AtomicBoolean(true);
    /**
     * 配置文件路径
     */
    private static final String CONFIG_FILE = "config/proxy.properties";
    private final ThreadPoolManager threadPoolManager;
    private final TrafficStats trafficStats;
    private final BackendHealthChecker healthChecker;
    /**
     * 负载均衡器
     */
    private final RoundRobinLoadBalancer loadBalancer;
    /**
     * 监控任务Future
     */
    private final Future<?> watcherFuture;

    public UdpProxyServer() throws SocketException {
        this.threadPoolManager = new ThreadPoolManager();
        this.trafficStats = new TrafficStats(threadPoolManager);
        this.healthChecker = new BackendHealthChecker(threadPoolManager.getScheduleExecutor());
        this.loadBalancer = new RoundRobinLoadBalancer(healthChecker);
        // 提交配置文件监控任务到线程池
        this.watcherFuture = threadPoolManager.getScheduleExecutor().submit(this::watchConfigFile);
    }

    public void start() throws InterruptedException {
        ProxyConfig config = ProxyConfig.getInstance();
        int port = config.getMulticastPort();
        String multicastAddress = config.getMulticastAddress();

        EventLoopGroup group = new NioEventLoopGroup();
        try {
            NetworkInterface ni = getNetworkInterface();
            Bootstrap bootstrap = new Bootstrap();
            bootstrap.group(group)
                    .channelFactory((ChannelFactory<NioDatagramChannel>) () -> new NioDatagramChannel(InternetProtocolFamily.IPv4))
                    .option(ChannelOption.IP_MULTICAST_TTL, 16)
                    .option(ChannelOption.IP_MULTICAST_IF, ni)
                    .option(ChannelOption.SO_REUSEADDR, true)
                    .option(ChannelOption.SO_BROADCAST, true)
                    .option(ChannelOption.SO_RCVBUF, 1024 * 1024 * 20)
                    .option(ChannelOption.RCVBUF_ALLOCATOR, new FixedRecvByteBufAllocator(65535))
                    .handler(new ChannelInitializer<DatagramChannel>() {
                        @Override
                        protected void initChannel(DatagramChannel ch) {
                            ch.pipeline().addLast(new UdpPacketHandler());
                        }
                    });
            ChannelFuture future = bootstrap.bind(port).sync();
            // 加入组播地址
            InetSocketAddress groupAddress = new InetSocketAddress(multicastAddress, port);
            ((NioDatagramChannel) future.channel()).joinGroup(groupAddress, ni).sync();
            logger.info("UDP代理服务器已启动，监听组播:{}和端口:{}", multicastAddress, port);

            // 启动流量统计
            trafficStats.start();

            // 启动健康检查
            healthChecker.start();

            future.channel().closeFuture().await();
        } catch (SocketException e) {
            throw new RuntimeException(e);
        } finally {
            group.shutdownGracefully();
        }
    }

    private class UdpPacketHandler extends SimpleChannelInboundHandler<DatagramPacket> {
        @Override
        protected void channelRead0(ChannelHandlerContext ctx, DatagramPacket packet) {
            if (!RUNNING.get()) {
                return;
            }
            // 转发数据包到后端服务器
            forwardPacket(ctx, packet);
        }

        private void forwardPacket(ChannelHandlerContext ctx, DatagramPacket receivedPacket) {
            Backend backend;
            try {
                // 使用健康检查器获取下一个后端
                backend = loadBalancer.getNextBackend();
                if (backend != null) {
                    // 创建转发数据包
                    ByteBuf content = receivedPacket.content();
                    byte[] data = ByteBufUtil.getBytes(content);
                    int size = data.length;
                    DatagramPacket forwardPacket = new DatagramPacket(
                            Unpooled.copiedBuffer(data),
                            new InetSocketAddress(backend.getHost(), backend.getPort())
                    );

                    // 发送转发数据包
                    ctx.writeAndFlush(forwardPacket);
                    trafficStats.recordForwarded(size);
                    logger.debug("已转发数据包到 {}:{}", backend.getHost(), backend.getPort());
                } else {
                    logger.error("没有可用的健康后端服务器进行转发");
                }
            } catch (Exception e) {
                logger.error("转发数据包失败", e);
            }
        }

        @Override
        public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) {
            logger.error("处理数据包时出错", cause);
            ctx.close();
        }
    }

    /**
     * ip正则表达式
     */
    private final static String IP_REGEX = "^(1\\d{2}|2[0-4]\\d|25[0-5]|[1-9]\\d|[1-9])\\."
            + "(1\\d{2}|2[0-4]\\d|25[0-5]|[1-9]\\d|\\d)\\."
            + "(1\\d{2}|2[0-4]\\d|25[0-5]|[1-9]\\d|\\d)\\."
            + "(1\\d{2}|2[0-4]\\d|25[0-5]|[1-9]\\d|\\d)$";

    /**
     * 判断是否是IP地址的字符串
     *
     * @param ipAddress IP地址
     * @return 是否是IP
     */
    public static boolean isIp(String ipAddress) {
        return ipAddress.matches(IP_REGEX);
    }

    /**
     * 获取网络接口
     *
     * @return {@link NetworkInterface}
     * @throws SocketException 套接字例外
     */
    private NetworkInterface getNetworkInterface() throws SocketException {
        Enumeration<NetworkInterface> netInterfaces = NetworkInterface.getNetworkInterfaces();
        while (netInterfaces.hasMoreElements()) {
            NetworkInterface ni = netInterfaces.nextElement();
            String displayName = ni.getDisplayName();
            String name = ni.getName();
            // 判断网络接口的名称不能为 lo获得docker 我们要找的应该是 eth0
            if (!name.contains("docker") && !name.contains("lo") && !displayName.contains("VMware") && !displayName.contains("virbr") && !displayName.contains("tun")) {
                Enumeration<InetAddress> ips = ni.getInetAddresses();
                while (ips.hasMoreElements()) {
                    InetAddress inetAddress = ips.nextElement();
                    // 不是回环地址 127.0.0.1
                    if (!inetAddress.isLoopbackAddress()) {
                        String ip = inetAddress.getHostAddress();
                        if (isIp(ip)) {
                            return ni;
                        }
                    }
                }
            }
        }
        return NetUtil.LOOPBACK_IF;
    }

    /**
     * 监控配置文件变化
     */
    private void watchConfigFile() {
        try {
            Path configPath = Paths.get(CONFIG_FILE);
            WatchService watchService = FileSystems.getDefault().newWatchService();
            Path dir = configPath.getParent();
            if (dir == null) {
                dir = Paths.get(".");
            }
            dir.register(watchService, StandardWatchEventKinds.ENTRY_MODIFY);

            while (RUNNING.get()) {
                WatchKey key = watchService.poll(1, TimeUnit.SECONDS);
                if (key == null) {
                    continue;
                }

                boolean needsReload = false;
                for (WatchEvent<?> event : key.pollEvents()) {
                    Path changed = (Path) event.context();
                    if (changed.toString().equals(configPath.getFileName().toString())) {
                        needsReload = true;
                        break;
                    }
                }

                key.reset();

                if (needsReload) {
                    threadPoolManager.getScheduleExecutor().schedule(this::reload, 500, TimeUnit.MILLISECONDS);
                }
            }
        } catch (Exception e) {
            if (RUNNING.get()) {
                logger.error("配置文件监控出错", e);
            }
        }
    }

    /**
     * 停止服务器
     */
    public void stop() {
        RUNNING.set(false);
        // 取消任务
        if (watcherFuture != null) {
            watcherFuture.cancel(true);
        }
        // 关闭线程池
        threadPoolManager.shutdown();
        // 停止流量统计
        trafficStats.stop();
        // 停止健康检查
        healthChecker.stop();
    }

    /**
     * 重新加载配置
     */
    private void reload() {
        logger.info("检测到配置文件变更，正在重新加载...");
        ProxyConfig.reload();
        healthChecker.resetHealth();
        logger.info("配置重新加载成功");
    }

    /**
     * 程序入口点
     *
     * @param args 参数
     * @throws SocketException      套接字异常
     */
    public static void main(String[] args) throws SocketException {
        UdpProxyServer server = new UdpProxyServer();

        Runtime.getRuntime().addShutdownHook(new Thread(server::stop));

        try {
            server.start();
            Thread.currentThread().join();
        } catch (Exception e) {
            logger.error("服务器启动出错", e);
            System.exit(1);
        }
    }
} 